import rsa
from rsa import PublicKey
import base64

def rsa_encrypt(text, public_key_str, segment_length=20):
    """
    使用RSA公钥对字符串进行分段加密，与JavaScript代码逻辑一致
    
    参数:
        text: 待加密的字符串
        public_key_str: RSA公钥字符串（PEM格式）
        segment_length: 每段加密的长度，默认为20
        
    返回:
        加密后用逗号分隔的字符串
    """
    # 处理公钥，添加必要的换行符
    public_key_str = public_key_str.replace(' ', '\n')
    public_key_str = f"-----BEGIN PUBLIC KEY-----\n{public_key_str}\n-----END PUBLIC KEY-----"
    
    # 导入公钥
    public_key = rsa.PublicKey.load_pkcs1_openssl_pem(public_key_str.encode())
    
    encrypted_segments = []
    total_length = len(text)
    start = 0
    
    while True:
        # 计算当前分段的结束位置
        end = start + segment_length
        if end > total_length:
            end = total_length
            
        # 截取当前分段
        segment = text[start:end]
        if not segment:
            break
            
        # 加密当前分段（RSA加密需要先将字符串转为字节）
        encrypted_bytes = rsa.encrypt(segment.encode('utf-8'), public_key)
        # 转为Base64编码的字符串
        encrypted_str = base64.b64encode(encrypted_bytes).decode('utf-8')
        
        encrypted_segments.append(encrypted_str)
        
        # 移动到下一段
        start += segment_length
        
        if start >= total_length:
            break
    
    return ','.join(encrypted_segments)

# 测试代码
if __name__ == "__main__":
    # 公钥（与JavaScript中的l变量对应）
    public_key = "MIGfMA0GCSqGSIb3DQEBAQUAA4GNADCBiQKBgQC3ZENUrDErqKZZCb2PtkFgi1De5wCrcEgd9CqoaAfvSGxV8V92MABFbS6+NXf2xJsgsP+fjMZljf9mBhy/jawRfAwlOQGqipu4cvtczUZzPcdE3DcjznjWDbhL66/QwJgYipcq8JU5mvJUxxUHhmshc6Twlfhw4me7F9V6lBpaowIDAQAB"
    
    # 待加密的文本
    test_text = r'{"current":5,"size":40,"monthId":"202507","stationId":"","stationName":"","cityCode":"0755","rentFree":"","powerFree":"","isSzdg":"","isHightEnergy":"","ossStationClassify":[],"stationState":[],"propertyName":[],"stationType":[],"stationClassifyName":[],"propChar":[],"contractNoGlobal":"","buildCode":"","towerCode":"","costTowerRepeated":"","isOssDelete":"","isLowP":""}'
    
    # 进行加密
    encrypted_result = rsa_encrypt(test_text, public_key)
    print("加密结果:", encrypted_result)
    